#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
修复HanLP与Python 3.12兼容性问题的脚本
"""

import os
import sys
import subprocess
import platform
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)
logger = logging.getLogger("hanlp-fix")

def print_step(step, message):
    """打印步骤信息"""
    print(f"\n\033[1;34m[步骤 {step}] {message}\033[0m")

def run_command(command):
    """运行命令并返回结果"""
    print(f"\033[90m执行: {command}\033[0m")
    result = subprocess.run(command, shell=True, text=True, capture_output=True)
    if result.returncode != 0:
        print(f"\033[91m命令执行失败: {result.stderr}\033[0m")
        return False
    print(f"\033[92m命令执行成功\033[0m")
    return True

def check_pip_version():
    """检查pip版本"""
    print_step(1, "检查pip版本")
    result = subprocess.run(
        [sys.executable, "-m", "pip", "--version"], 
        text=True, 
        capture_output=True
    )
    print(f"当前pip版本: {result.stdout.strip()}")
    return run_command(f"{sys.executable} -m pip install --upgrade pip")

def check_python_version():
    """检查Python版本"""
    print_step(2, "检查Python版本")
    version = sys.version.split()[0]
    print(f"当前Python版本: {version}")
    
    # 检查是否为Python 3.12
    if version.startswith("3.12"):
        print("\033[93m警告: 检测到Python 3.12，部分HanLP依赖可能存在兼容性问题\033[0m")
        print("建议使用Python 3.8-3.11版本以获得最佳兼容性")
        return True
    return False

def install_setuptools():
    """安装或更新setuptools"""
    print_step(3, "更新setuptools")
    return run_command(f"{sys.executable} -m pip install --upgrade setuptools wheel")

def install_numpy():
    """单独安装numpy"""
    print_step(4, "安装numpy")
    return run_command(f"{sys.executable} -m pip install 'numpy>=1.22.0,<1.25.0'")

def install_torch():
    """安装PyTorch"""
    print_step(5, "安装PyTorch")
    
    # 判断操作系统和CUDA支持
    if platform.system() == "Windows":
        return run_command(f"{sys.executable} -m pip install torch torchvision torchaudio")
    else:  # Linux/Mac
        return run_command(f"{sys.executable} -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu")

def install_transformers():
    """安装transformers"""
    print_step(6, "安装transformers")
    return run_command(f"{sys.executable} -m pip install 'transformers>=4.26.0,<4.30.0'")

def install_hanlp():
    """安装HanLP"""
    print_step(7, "安装HanLP")
    return run_command(f"{sys.executable} -m pip install 'hanlp>=2.1.0b36'")

def install_flask():
    """安装Flask"""
    print_step(8, "安装Flask")
    return run_command(f"{sys.executable} -m pip install flask==2.3.3 flask-cors==3.0.10 werkzeug==2.3.7")

def main():
    """主函数"""
    print("\033[1;32m===== HanLP依赖修复工具 =====\033[0m")
    
    # 检查pip版本
    if not check_pip_version():
        sys.exit(1)
    
    # 检查Python版本
    is_py312 = check_python_version()
    
    # 安装setuptools
    if not install_setuptools():
        sys.exit(1)
    
    # 如果是Python 3.12，单独安装各个依赖
    if is_py312:
        print("\033[93m检测到Python 3.12，将单独安装各个依赖以解决兼容性问题\033[0m")
        
        # 安装numpy
        if not install_numpy():
            sys.exit(1)
        
        # 安装PyTorch
        if not install_torch():
            sys.exit(1)
        
        # 安装transformers
        if not install_transformers():
            sys.exit(1)
        
        # 安装HanLP
        if not install_hanlp():
            sys.exit(1)
        
        # 安装Flask
        if not install_flask():
            sys.exit(1)
    else:
        # 使用requirements.txt安装
        print_step(3, "使用requirements.txt安装依赖")
        if not run_command(f"{sys.executable} -m pip install -r requirements.txt"):
            sys.exit(1)
    
    print("\n\033[1;32m===== 依赖安装完成 =====\033[0m")
    print("如果安装过程中仍然遇到问题，建议尝试以下方案：")
    print("1. 使用Python 3.8-3.11版本")
    print("2. 创建并使用虚拟环境")
    print("3. 手动安装各个依赖，特别是numpy和torch")

if __name__ == "__main__":
    main() 